import os
import cv2
import re

import numpy as np


def extract_prefix(filename):
    pattern = re.compile(r'(.*)tr\d+-tc\d+\.png')
    match = pattern.match(filename)
    if match:
        return match.group(1)
    return ""


def split_and_save_small_tiles(input_dir, output_dir_256, output_dir_512):
    os.makedirs(output_dir_256, exist_ok=True)
    os.makedirs(output_dir_512, exist_ok=True)

    image_files = sorted([f for f in os.listdir(input_dir) if f.endswith('.png')])

    # 计算最大行和列编号
    max_row = 0
    max_col = 0
    for image_file in image_files:
        match = re.search(r'tr(\d+)-tc(\d+)', image_file)
        if not match:
            continue
        tr = int(match.group(1))
        tc = int(match.group(2))
        max_row = max(max_row, tr)
        max_col = max(max_col, tc)

    for image_file in image_files:
        prefix = extract_prefix(image_file)
        match = re.search(r'tr(\d+)-tc(\d+)', image_file)
        if not match:
            continue

        tr = int(match.group(1))
        tc = int(match.group(2))

        image_path = os.path.join(input_dir, image_file)
        image = cv2.imread(image_path)
        if image is None:
            raise FileNotFoundError(f"Unable to open image file: {image_path}")

        tile_size = 256
        rows, cols = image.shape[0] // tile_size, image.shape[1] // tile_size

        for i in range(rows):
            for j in range(cols):
                small_tile = image[i * tile_size:(i + 1) * tile_size, j * tile_size:(j + 1) * tile_size]

                new_tr = (tr - 1) * 4 + i + 1
                new_tc = (tc - 1) * 4 + j + 1

                small_tile_filename = f"{prefix}tr{new_tr}-tc{new_tc}.png"
                small_tile_path = os.path.join(output_dir_256, small_tile_filename)
                cv2.imwrite(small_tile_path, small_tile)
                print(f"Saved {small_tile_filename} to {small_tile_path}")

    # 生成512x512图像
    for row in range(1, max_row * 4 - 1, 1):  # 从1到max_row*4，步长为2
        for col in range(1, max_col * 4 - 1, 1):  # 从1到max_col*4，步长为2
            tile_paths = [
                os.path.join(output_dir_256, f"{prefix}tr{row}-tc{col}.png"),
                os.path.join(output_dir_256, f"{prefix}tr{row}-tc{col + 1}.png"),
                os.path.join(output_dir_256, f"{prefix}tr{row + 1}-tc{col}.png"),
                os.path.join(output_dir_256, f"{prefix}tr{row + 1}-tc{col + 1}.png")
            ]

            if all(os.path.exists(path) for path in tile_paths):
                tiles = [cv2.imread(path) for path in tile_paths]
                large_tile = np.zeros((2 * tile_size, 2 * tile_size, 3), dtype=np.uint8)

                large_tile[0:tile_size, 0:tile_size] = tiles[0]
                large_tile[0:tile_size, tile_size:2 * tile_size] = tiles[1]
                large_tile[tile_size:2 * tile_size, 0:tile_size] = tiles[2]
                large_tile[tile_size:2 * tile_size, tile_size:2 * tile_size] = tiles[3]

                large_tile_filename = f"{prefix}tr{row}-tc{col}.png"
                large_tile_path = os.path.join(output_dir_512, large_tile_filename)
                cv2.imwrite(large_tile_path, large_tile)
                print(f"Saved {large_tile_filename} to {large_tile_path}")


# 示例用法
input_dir = 'E:\\wafer52\\11867_32nm_coarse_align2(1)'  # 替换为实际输入目录路径
output_dir_256 = 'E:\\wafer52\\11867_32nm_256_split(1)'  # 替换为实际256x256小图输出目录路径
output_dir_512 = 'E:\\wafer52\\11867_32nm_512_combine(1)'  # 替换为实际512x512小图输出目录路径

split_and_save_small_tiles(input_dir, output_dir_256, output_dir_512)
